'''
Created on 2011-06-21

@author: Andrew Roth
'''
from collections import defaultdict, namedtuple

'''
Pure python base counter. Kept in tree for performance comparison.
'''


class BaseCounter(object):
    row = namedtuple('BaseCountRow', ['ref', 'position', 'counts'])
    ascii_offset = 33
    
    def __init__(self, bam_file, min_base_qual=10, min_map_qual=10):
        self._bam_file = bam_file
        
        self.min_base_qual = min_base_qual
        self.min_map_qual = min_map_qual
        
        self.refs = bam_file.references
    
    def iter_ref(self, ref):
        for pileup_column in self._bam_file.pileup(ref):
            counts_dict = self._parse_pileup_column(pileup_column)
            
            counts = {
                      'A' : counts_dict['A'],
                      'C' : counts_dict['C'],
                      'G' : counts_dict['G'],
                      'T' : counts_dict['T']
                      }
            
            position = pileup_column.pos + 1
            
            row = BaseCounter.row(ref, position, counts)
            
            yield row

    def _parse_pileup_column(self, pileup_column):
        base_counts = defaultdict(int)
        
        for read in pileup_column.pileups:
            if read.is_del:
                continue
            
            qpos = read.qpos
                 
            map_qual = read.alignment.mapq
            
            if map_qual < self.min_map_qual:
                continue            
            
            base_qual = ord(read.alignment.qual[qpos]) - BaseCounter.ascii_offset            
            
            if base_qual < self.min_base_qual:
                continue
            
            base = read.alignment.seq[qpos]
            base_counts[base] += 1
        
        return base_counts

